Skip to content

Correction: the jax-compile-stall is not a compile stall - #279

Merged
Jammy2211 merged 2 commits into
mainfrom
claude/jax-vmap-jit-stall-swz2tc
Aug 23, 2026
Merged

Correction: the jax-compile-stall is not a compile stall#279
Jammy2211 merged 2 commits into
mainfrom
claude/jax-vmap-jit-stall-swz2tc

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Follow-up to #278. The two 1800s re-timing runs that were still in flight at close-out finished, and the faulthandler stack they captured contradicts the premise the whole epic was filed under.

The evidence

Both runs dumped at 1440s, exactly as PyAutoFit#1518 intended. Two different repos, two different scripts, the same stack:

Thread 0x00007fbf1f0e4b80 (most recent call first):
  File ".../jax/_src/api.py", line 2764 in try_to_block
  File ".../jax/_src/api.py", line 2781 in block_until_ready
  File ".../autofit/non_linear/jax_compile.py", line 264 in wrapper
  File ".../scripts/imaging/jax_likelihood/mge_group.py", line 175 in <module>

jax_compile.py:264 is the jax.block_until_ready(result) call — the second half of the wrapped first call. The process is not tracing, lowering or compiling. func(*args, **kwargs) returned. What never returns is the wait for the result to materialize.

What this changes

  • Every marker, issue and prompt in this epic calls it an "intermittent XLA compile stall". That name is wrong, and has been since the first quarantine on 2026-08-01 — it was a guess made before any evidence existed. Corrected in the record, epics.md, the campaign ledger, and the follow-up smoke-timing task, so the speed-up work doesn't restart from the wrong question.
  • A defect in log_on_first_compile, worth fixing before anyone reads another of these logs: the heartbeat prints JAX jit still compiling ... 1770s elapsed while the process sits in block_until_ready. The wrapper can't tell which half it's in and says "compiling" regardless. That wording actively reinforced the wrong diagnosis for the full 1770 seconds.
  • Both entries are AMBIGUOUS even at 1800s — 2/2 capped on both legs. With the 300s round that's 20 consecutive cap hits and zero completions for each, while multi_dataset/mge.py's own marker records it finishing in 32s standalone.

Where a resumption should start

Not "why is XLA slow to compile" but "why does block_until_ready never return". The record lists three candidates: a device-transfer / async-dispatch hang (try_to_block is exactly where a never-arriving buffer parks); the compilation cache, still live but for a different reason (a cache read satisfying the compile instantly and leaving execution to hang fits what's seen); and the vmap(jit) ordering result reinterpreted — the ordering changes the shape of the executed computation, which is a more plausible route to an execution hang than to a compile one.

Note

The epic stays closed as partial — this sharpens the diagnosis, it does not root-cause it, and nothing is un-quarantined. Branch restarted from main after #278 merged, per the merged-PR rule.

lifecycle.py check: OK. Index and dashboard regenerated.


Generated by Claude Code

James Nightingale and others added 2 commits August 23, 2026 23:09
The two 1800s runs dumped faulthandler tracebacks at 1440s. Both repos, both
scripts, the same stack: parked in jax.block_until_ready / try_to_block — the
EXECUTION half of the first call. func(*args, **kwargs) returned; what never
returns is the wait for the result to materialize.

So the epic's name, and every marker calling this an intermittent XLA compile
stall, inherit a guess made before there was evidence. Corrected in the record,
the epic entry, the ledger and the follow-up task so the smoke speed-up work
does not restart from the wrong question.

Also recorded: the heartbeat says 'still compiling' while the process sits in
block_until_ready — log_on_first_compile cannot tell which half it is in. Fix
the wording before anyone reads another of these logs.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
The refresh check on PR #279 was right: my PyAutoBrain checkout was one commit
behind main (#260, which changed the board logo markup), so my render differed
from the one CI produces. Regenerated against main.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015qk7hoavMnFyPtW4toYn8K
@Jammy2211
Jammy2211 merged commit 9c1aafa into main Aug 23, 2026
2 checks passed
@github-actions
github-actions Bot deleted the claude/jax-vmap-jit-stall-swz2tc branch August 25, 2026 14:02
Jammy2211 pushed a commit that referenced this pull request Aug 26, 2026
…oArray mapper bug

Reproduced the failure bit-for-bit with a runnable stack built in-session (both
Python legs, source-installed libraries at the retime run's exact versions), and
the prompt's own UPDATE 2026-08-24 diagnosis does not survive it.

Refuted, with evidence:
  1. NOT "Python 3.13 only" — 3.12 fails identically on the same host, so the CI
     split was runner hardware, not CPython.
  2. NOT a PDIP branch flip — pdip_iter is identical eager vs jit, and the gap is
     bit-identical across 5 (nnls_solver_tol, nnls_max_iter) policies.
  3. Therefore pinning solver policy cannot fix it; that direction is dead.

Actual cause: a discrete bilinear cell-assignment flip in PyAutoArray's
rectangular mapper. transform() ends in clip(F_q, 0, 1), so saturated points land
on exactly-integer indices where ix_up = ceil(g) collapses onto ix_down; one ULP
in the traced grid then jumps a point's weight a whole mesh row. Underneath it the
row weights are mirrored — the current code fails a linear-reproduction test by
~a full cell in the row axis while the column axis is exact.

Spawns draft/bug/autoarray/rectangular_mapper_bilinear_row_weights.md as the real
fix (human-required: it changes reconstructions library-wide and 16 workspace
scripts' hardcoded constants). The NEEDS_FIX park stays until that lands.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016azuS2UbS3mFvfxkDFGsKj
Jammy2211 pushed a commit that referenced this pull request Aug 26, 2026
…#279 / PyAutoArray#490)

draft/test/workspaces/physical_model_check_when_speeding_up_smoke.md
  A speed-up on a validation script is not done until that script's lens model is
  known to be physical. #490 survived 11 months partly because the jax_likelihood
  scripts assert at deliberately-wrong models: rectangular_rtu.py pins the mass
  centre to (0.3,-0.3) against a (0,0) truth, and the jax_test dataset carries
  Sersic+Exponential lens light the mass-only models have no component for (there,
  truth measurably fits WORSE than the offset model). Records the pixelized-source
  caveat: assert on log_likelihood and source recovery, not the figure of merit,
  which for a pixelization is the evidence and moves for unrelated reasons.

draft/test/autoarray/final_numerics_audit_of_every_mesh_interpolator.md
  One systematic pass over every mesh interpolator using the five checks that
  actually discriminated in #490 — linear reproduction, continuity across integer
  boundaries, convergence under refinement, ground-truth source recovery, and a
  bit-identical known-good control. Carries the specific suspicions: saturating
  transforms feeding a discretisation, corner/weight pairing, n_knots not scaling
  with mesh_pixels (roundtrip drift 0.0055 -> 0.0238 -> 0.101 index units at
  n=16/32/64), and guard-node conventions. Requires every new test to be shown
  failing against a broken variant.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016azuS2UbS3mFvfxkDFGsKj
Jammy2211 pushed a commit that referenced this pull request Aug 26, 2026
…regression, not an NNLS branch flip

Both PRs merged library-first and verified as ancestors of their main with 0
unmerged commits:
  PyAutoArray#490            merged 158db384
  autolens_workspace_test#280 merged 8be3d598

The filed prompt's own diagnosis was refuted on three counts by building both
Python legs from source and reproducing the failure bit-for-bit: not 3.13-only
(3.12 fails identically on the same host — the CI split was runner hardware),
not a PDIP branch flip (pdip_iter identical eager vs jit; the gap bit-identical
across five solver policies), and therefore not fixable by pinning solver policy.

Actual cause: mirrored bilinear row weights plus a round-off-dependent cell
assignment in PyAutoArray's adaptive rectangular mapper, live since 2025-09-23.
Both are regressions — the mapper was correct at introduction (2025-06-24) and
the correct formulation still ships for the uniform mesh at
interpolator/rectangular_uniform.py:72-99, which is what the fix restores.

Moves active/ -> complete/2026/08/, releases the active.md claim, refreshes
complete/index.md and regenerates the dashboard in this commit.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016azuS2UbS3mFvfxkDFGsKj
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant